﻿//
// Created by wzh on 2020/12/29.
//

#include "splinefitter.h"


SplineFitter::SplineFitter() : m_bound_type(First), m_bound_l(0.0), m_bound_r(0.0) {};

void SplineFitter::setBoundary(BoundValueType type, double left, double right) {
    m_bound_type = type;
    m_bound_l = left;
    m_bound_r = right;
};

bool SplineFitter::setPoints(const Vecd &x, const Vecd &y) {

    if (x.size() != y.size() && x.size() < 3) {
        return false;
    }

    size_t n = x.size();
    m_x = x;
    m_y = y;

    bubbleSort(m_x.data(), m_y.data(), n);
    auto D = m_y;
    Matd A(n, 3);
    A(0, 0) = 0.0;
    A(n - 1, 2) = 0.0;
    A.col(1).fill(2.0);
    Arrd H;
    H = m_x.tail(n - 1) - m_x.head(n - 1);
    for (size_t i = 1; i < n - 1; ++i) {
        auto temp = H[i - 1] + H[i];
        A(i, 2) = H[i] / temp;   //λ
        A(i, 0) = 1 - A(i, 2);      //μ
    }

    //计算差商D
    Arrd y1 = (D.tail(n - 1) - D.head(n - 1)) / H;
    D.segment(1, n - 2) = 6 * (y1.tail(n - 2) - y1.head(n - 2)) / (m_x.tail(n - 2) - m_x.head(n - 2));
    //边界条件
    switch (m_bound_type) {
        case First:
            A(0, 2) = 0;
            A(n - 1, 0) = 0;
            D[0] = 0;
            D[n - 1] = 0;
            break;
        case Second:
            A(0, 2) = 1;
            A(n - 1, 0) = 1;
            D[0] = 6.0 / H[0] * (y1[0] - m_bound_l);
            D[n - 1] = 6.0 / H[n - 2] * (m_bound_r - y1[n - 2]);
            break;
        case Third:
            A(0, 2) = -2;
            A(n - 1, 0) = -2;
            //X2不是12的原因是d本来已经是六倍差商了
            D[0] = -2.0 * H[0] * (D[1] - D[2]) / (m_x[0] - m_x(3));
            D[n - 1] = 2.0 * H[n - 2] * (D[n - 3] - D[n - 2]) / (m_x[n - 4] - m_x(n - 1));
            break;
    }

    m_2 = BandSolver::solveCompressed(1, 1, A, D) / 2;
    m_3.resize(n);
    m_3.head(n - 1) = (m_2.tail(n - 1) - m_2.head(n - 1)) / (3 * H);
    m_1.resize(n);
    m_1.head(n - 1) = y1 - (H / 3) * (2 * m_2.head(n - 1) + m_2.tail(n - 1));
    m_3[n - 1] = 0.0;
    m_1[n - 1] = 3.0 * m_3[n - 2] * H[n - 2] * H[n - 2] + 2.0 * m_2[n - 2] * H[n - 2] + m_1[n - 2];

    return true;
}


double SplineFitter::operator()(double x) {
    auto n = m_x.size();
    // find the closest point m_x[idx] < x, idx=0 even if x<m_x[0]
    auto it = std::lower_bound(m_x.data(), m_x.data() + n, x);
    int idx = std::max(int(it - m_x.data()) - 1, 0);
    double h = x - m_x[idx];
    double ret;
    if (x < m_x[0]) {
        // extrapolation to the left
        ret = (m_2[0] * h + m_3[0]) * h + m_y[0];
    } else if (x > m_x[n - 1]) {
        // extrapolation to the right
        ret = (m_2[n - 1] * h + m_3[n - 1]) * h + m_y[n - 1];
    } else {
        // interpolation
        ret = ((m_3[idx] * h + m_2[idx]) * h + m_1[idx]) * h + m_y[idx];
    }
    return ret;
}


